/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file edge32_hor_c3.cpp
 */
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelEdge32HorC3 {
public:
    __aicore__ inline KernelEdge32HorC3() {}
    __aicore__ inline void Init(GM_ADDR img, GM_ADDR out,
                                uint32_t singleHeigh1, uint32_t width, uint32_t heigh,
                                TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        if (GetBlockIdx() < 8) 
        {
            this->singleHeigh = 262;
            imgGm.SetGlobalBuffer((__gm__ uint8_t*)img + this->singleHeigh * GetBlockIdx() * width * 3, this->singleHeigh * width * 3);
            outGm.SetGlobalBuffer((__gm__ uint8_t*)out + this->singleHeigh * GetBlockIdx() * width, this->singleHeigh * width * 3);
        }
        else
        {
            this->singleHeigh = 192;
            imgGm.SetGlobalBuffer((__gm__ uint8_t*)img + 262*8*width * 3 + this->singleHeigh * (GetBlockIdx()-8) * width * 3, this->singleHeigh * width * 3);
            outGm.SetGlobalBuffer((__gm__ uint8_t*)out + 262*8*width + this->singleHeigh * (GetBlockIdx()-8) * width, this->singleHeigh * width * 3);
        }
        if(GetBlockIdx() == 14)
        {
            this->singleHeigh -= 2;
        }
        this->width = width;
        this->heigh = heigh;

        this->tileDataNum = 4864;
        pipeIn->InitBuffer(inQueueImg1, BUFFER_NUM, this->tileDataNum * 3 * sizeof(uint8_t));
        pipeIn->InitBuffer(inQueueImg2, BUFFER_NUM, this->tileDataNum * 3 * sizeof(uint8_t));
        pipeIn->InitBuffer(outQueueOut, BUFFER_NUM, this->tileDataNum * sizeof(uint8_t));

        pipeIn->InitBuffer(QueueTmp1, this->tileDataNum * 6 * sizeof(half));
        pipeIn->InitBuffer(QueueTmp2, this->tileDataNum * 6 * sizeof(half));
    }
    __aicore__ inline void Process() {
        GmInit();
        int32_t loopCount = this->singleHeigh;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }
private:
    __aicore__ inline void GmInit() {   
        LocalTensor<uint8_t> outLocal = outQueueOut.AllocTensor<uint8_t>();
        Duplicate(outLocal.ReinterpretCast<int16_t>(), (int16_t)0, (4896+32)/2);
        outQueueOut.EnQue<uint8_t>(outLocal);

        outLocal = outQueueOut.DeQue<uint8_t>();
        if(GetBlockIdx() == 0)
        {
            DataCopy(outGm, outLocal, 4896+32);
            PipeBarrier<PIPE_MTE3>();
            DataCopy(outGm[4887*(3440-1)], outLocal, 4896);
        }
        for(int32_t i = 0; i < this->singleHeigh; i++)
        {
            DataCopy(outGm[(i * this->width + this->width + this->width-16)], outLocal, 32);
        }
        
        outQueueOut.FreeTensor(outLocal);
    }
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<uint8_t> img1Local = inQueueImg1.AllocTensor<uint8_t>();
        LocalTensor<uint8_t> img2Local = inQueueImg2.AllocTensor<uint8_t>();
        DataCopy(img1Local, imgGm[(progress * this->width+16) * 3], this->tileDataNum *3);
        DataCopy(img2Local, imgGm[(progress * this->width+this->width*2+16) * 3], this->tileDataNum *3);
        inQueueImg1.EnQue(img1Local);
        inQueueImg2.EnQue(img2Local);
    }
    __aicore__ inline void Compute(int32_t progress) {
        LocalTensor<uint8_t> img1Local = inQueueImg1.DeQue<uint8_t>();
        LocalTensor<uint8_t> img2Local = inQueueImg2.DeQue<uint8_t>();
        LocalTensor<uint8_t> outLocal = outQueueOut.AllocTensor<uint8_t>();
        auto tmp1 = QueueTmp1.Get<half>();
        auto tmp2 = QueueTmp2.Get<half>();
        uint32_t mask = this->tileDataNum;
        uint64_t rsvdCnt = 0;
        Cast(tmp1, img1Local, RoundMode::CAST_NONE, this->tileDataNum*3);
        Cast(tmp2, img2Local, RoundMode::CAST_NONE, this->tileDataNum*3);
        Sub(tmp2, tmp2, tmp1, this->tileDataNum*3);
        Abs(tmp2, tmp2, this->tileDataNum*3);
        uint32_t length = this->width - 16 - 16;
        TransposeParamsExt transposeParams;
        transposeParams.nSize = 2;
        transposeParams.cSize = 3;
        transposeParams.hSize = 152;
        transposeParams.wSize = 16;
        transposeParams.transposeType = TransposeType::TRANSPOSE_NHWC2NCHW;
        Transpose(tmp1, tmp2, img1Local, transposeParams);
        Duplicate(tmp2[4855], (half)0, 9);
        Max(tmp2, tmp1, tmp1[2432], 2432);
        Max(tmp2, tmp2, tmp1[2432*2], 2432);
        Max(tmp2[2432], tmp1[2432*3], tmp1[2432*4], 2423);
        Max(tmp2[2432], tmp2[2432], tmp1[2432*5], 2423);
        Cast(outLocal, tmp2, RoundMode::CAST_NONE, 4864);
        inQueueImg1.FreeTensor(img1Local);
        inQueueImg2.FreeTensor(img2Local);
        outQueueOut.EnQue<uint8_t>(outLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<uint8_t> outLocal = outQueueOut.DeQue<uint8_t>();
        DataCopy(outGm[(progress * this->width + this->width + 16)], outLocal, 4864);
        outQueueOut.FreeTensor(outLocal);
    }

private:
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueImg1, inQueueImg2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOut;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2, QueuePattern;

    GlobalTensor<uint8_t> imgGm;
    GlobalTensor<uint8_t> outGm;
    GlobalTensor<uint16_t> outGm_16;

    uint32_t singleHeigh;
    uint32_t tileDataNum;
    uint32_t width;
    uint32_t heigh;
};
extern "C" __global__ __aicore__ void edge32_hor_c3(GM_ADDR img, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    KernelEdge32HorC3 op;
    op.Init(img, out,
            tiling_data.singleHeigh, tiling_data.width, tiling_data.heigh, 
            &pipe);  
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_VECTOR_CORE);
    op.Process();
}